from typing import Any

import torch
import torch_npu
from pytorch_lightning import LightningModule
from torch_npu.utils.tensor_methods import torch_device_guard
from torch.nn import Module

import pytorch_lightning


@torch_device_guard
def to(self, *args: Any, **kwargs: Any):

    if args and isinstance(args[0], str) and 'npu' in args[0]:
        args = tuple([list(args)[0].replace('npu', torch_npu.npu.native_device)])
    if kwargs and 'npu' in str(kwargs.get("device", "")):
        kwargs['device'] = kwargs['device'].replace("npu", torch_npu.npu.native_device)

    out = torch._C._nn._parse_to(*args, **kwargs)
    self._DeviceDtypeModuleMixin__update_properties(device=out[0], dtype=out[1])
    return Module.to(self, *args, **kwargs)

def apply_module_patch():
    pytorch_lightning.core.LightningModule.to = to